#include "link.h"
#include "cpputil/cpputil/frexception.h"

Section::Section() : _addr(0), _sectionType(SECTION_TEXT), _sizeType(SIZE_16)
{
}

Section::Section(std::string name, uint64_t addr, SectionType sectionType, SizeType sizeType, BinBuffer data) : _name(name), _addr(addr), _sectionType(sectionType), _sizeType(sizeType), _data(data)
{
}

std::string Section::name()
{
	return _name;
}

uint64_t& Section::addr()
{
	return _addr;
}

SectionType Section::sectionType()
{
	return _sectionType;
}

SizeType Section::sizeType()
{
	return _sizeType;
}

BinBuffer& Section::data()
{
	return _data;
}

Symbol::Symbol() : _sectionOffset(0), _fileId(0)
{
}

Symbol::Symbol(uint64_t fileId, std::string name, std::string sectionName, uint64_t sectionOffset) : _fileId(fileId), _name(name), _sectionName(sectionName), _sectionOffset(sectionOffset)
{
}

uint64_t Symbol::fileId()
{
	return _fileId;
}

std::string Symbol::name()
{
	return _name;
}

std::string Symbol::sectionName()
{
	return _sectionName;
}

uint64_t Symbol::sectionOffset()
{
	return _sectionOffset;
}

Relocation::Relocation(std::string symbolName, std::string dstSectionName, uint64_t dstSectionOffset, bool isRelative) : _symbolName(symbolName), _dstSectionName(dstSectionName), _dstSectionOffset(dstSectionOffset), _isRelative(isRelative)
{
}

std::string Relocation::symbolName()
{
	return _symbolName;
}

std::string Relocation::dstSectionName()
{
	return _dstSectionName;
}

uint64_t Relocation::dstSectionOffset()
{
	return _dstSectionOffset;
}

bool Relocation::isRelative()
{
	return _isRelative;
}

bool Linker::findSymbol(ArrayMap<Symbol> symbols, std::string name, Symbol& res)
{
	for (auto& kv : symbols.data) {
		for (auto& symbol : kv.second) {
			if (symbol.name() == name) {
				res = symbol;
				return true;
			}
		}
	}
	return false;
}

bool Linker::findSymbol(ArrayMap<Symbol> symbols, uint64_t fileId, std::string name, Symbol& res)
{
	for (auto& symbol : symbols.data[fileId]) {
		if (symbol.name() == name) {
			res = symbol;
			return true;
		}
	}
	return false;
}

Linker::Linker() : fileId(0), _startAddr(0)
{
}

Linker& Linker::addSection(uint64_t fileId, Section section)
{
	sections.data[fileId].push_back(section);
	return *this;
}

Linker& Linker::addInSymbol(uint64_t fileId, Symbol symbol)
{
	in_symbols.data[fileId].push_back(symbol);
	return *this;
}

Linker& Linker::addOutSymbol(uint64_t fileId, Symbol symbol)
{
	out_symbols.data[fileId].push_back(symbol);
	return *this;
}

Linker& Linker::addRelocation(uint64_t fileId, Relocation relocation)
{
	relocations.data[fileId].push_back(relocation);
	return *this;
}

Linker& Linker::link(uint64_t baseAddr, std::string start_symbol)
{
	uint64_t sumAddr = 0;
	for (auto& kv : sections.data) {
		for (auto& section : kv.second) {
			section.addr() = sumAddr + baseAddr;
			sumAddr += section.data().length();
		}
	}
	for (auto& kv : relocations.data) {
		for (auto& relocation : kv.second) {
			Symbol resSymbol;
			bool findSection1 = false;
			if (findSymbol(in_symbols, kv.first, relocation.symbolName(), resSymbol)) {
				for (auto& section : sections.data[kv.first]) {
					if (section.name() == relocation.dstSectionName()) {
						findSection1 = true;
						if (relocation.isRelative()) {
							bool findSection2 = false;
							for (auto& symbolSection : sections.data[resSymbol.fileId()]) {
								if (symbolSection.name() == resSymbol.sectionName()) {
									findSection2 = true;
									switch (section.sizeType()) {
									case SIZE_16:
										*(uint16_t*)(&section.data()[relocation.dstSectionOffset()]) = uint16_t(symbolSection.addr() + resSymbol.sectionOffset() - 2 - relocation.dstSectionOffset() - section.addr());
										break;
									case SIZE_32:
										*(uint32_t*)(&section.data()[relocation.dstSectionOffset()]) = uint32_t(symbolSection.addr() + resSymbol.sectionOffset() - 4 - relocation.dstSectionOffset() - section.addr());
										break;
									case SIZE_64:
										*(uint64_t*)(&section.data()[relocation.dstSectionOffset()]) = uint64_t(symbolSection.addr() + resSymbol.sectionOffset() - 8 - relocation.dstSectionOffset() - section.addr());
										break;
									}
									break;
								}
							}
							if (!findSection2) {
								throw SectionNotFoundException(resSymbol.sectionName());
							}
						}
						else {
							bool findSection2 = false;
							for (auto& symbolSection : sections.data[resSymbol.fileId()]) {
								if (symbolSection.name() == resSymbol.sectionName()) {
									findSection2 = true;
									switch (section.sizeType()) {
									case SIZE_16:
										*(uint16_t*)(&section.data()[relocation.dstSectionOffset()]) += uint16_t(symbolSection.addr() + resSymbol.sectionOffset());
										break;
									case SIZE_32:
										*(uint32_t*)(&section.data()[relocation.dstSectionOffset()]) += uint32_t(symbolSection.addr() + resSymbol.sectionOffset());
										break;
									case SIZE_64:
										*(uint64_t*)(&section.data()[relocation.dstSectionOffset()]) += uint64_t(symbolSection.addr() + resSymbol.sectionOffset());
										break;
									}
									break;
								}
							}
							if (!findSection2) {
								throw SectionNotFoundException(resSymbol.sectionName());
							}
						}
						break;
					}
				}
				if (!findSection1) {
					throw SectionNotFoundException(relocation.dstSectionName());
				}
			}
			else if (findSymbol(out_symbols, relocation.symbolName(), resSymbol)) {
				findSection1 = false;
				for (auto& resSection : sections.data[kv.first]) {
					if (resSection.name() == relocation.dstSectionName()) {
						findSection1 = true;
						if (relocation.isRelative()) {
							bool findSection2 = false;
							for (auto& symbolSection : sections.data[resSymbol.fileId()]) {
								if (symbolSection.name() == resSymbol.sectionName()) {
									findSection2 = true;
									switch (resSection.sizeType()) {
									case SIZE_16:
										*(uint16_t*)(&(resSection.data()[relocation.dstSectionOffset()])) = uint16_t(symbolSection.addr() + resSymbol.sectionOffset() - 2 - relocation.dstSectionOffset() - resSection.addr());
										break;
									case SIZE_32:
										*(uint32_t*)(&resSection.data()[relocation.dstSectionOffset()]) = uint32_t(symbolSection.addr() + resSymbol.sectionOffset() - 4 - relocation.dstSectionOffset() - resSection.addr());
										break;
									case SIZE_64:
										*(uint64_t*)(&resSection.data()[relocation.dstSectionOffset()]) = uint64_t(symbolSection.addr() + resSymbol.sectionOffset() - 8 - relocation.dstSectionOffset() - resSection.addr());
										break;
									}
									break;
								}
							}
							if (!findSection2) {
								throw SectionNotFoundException(resSymbol.sectionName());
							}
						}
						else {
							bool findSection2 = false;
							for (auto& symbolSection : sections.data[resSymbol.fileId()]) {
								if (symbolSection.name() == resSymbol.sectionName()) {
									findSection2 = true;
									switch (resSection.sizeType()) {
									case SIZE_16:
										*(uint16_t*)(&resSection.data()[relocation.dstSectionOffset()]) += uint16_t(symbolSection.addr() + resSymbol.sectionOffset());
										break;
									case SIZE_32:
										*(uint32_t*)(&resSection.data()[relocation.dstSectionOffset()]) += uint32_t(symbolSection.addr() + resSymbol.sectionOffset());
										break;
									case SIZE_64:
										*(uint64_t*)(&resSection.data()[relocation.dstSectionOffset()]) += uint64_t(symbolSection.addr() + resSymbol.sectionOffset());
										break;
									}
									break;
								}
							}
							if (!findSection2) {
								throw SectionNotFoundException(resSymbol.sectionName());
							}
						}
						break;
					}
				}
				if (!findSection1) {
					throw SectionNotFoundException(relocation.dstSectionName());
				}
			}
			else {
				throw SymbolNotFoundException(relocation.symbolName());
			}
		}
	}
	if (start_symbol.size()) {
		for (auto& kv : out_symbols.data) {
			for (auto& symbol : kv.second) {
				if (symbol.name() == start_symbol) {
					for (auto& resSection : sections.data[kv.first]) {
						if (resSection.name() == symbol.sectionName()) {
							_startAddr = resSection.addr() + symbol.sectionOffset();
							return *this;
						}
					}
				}
			}
		}
	}
	_startAddr = baseAddr;
	return *this;
}

Linker& Linker::makeBin(std::string filename)
{
	File file(filename, "wb");
	for (auto& kv : sections.data) {
		for (auto& section : kv.second) {
			if (section.sectionType() == SECTION_BSS) {
				for (uint64_t i = 0; i < section.data().length(); i++) {
					fputc(0, file);
				}
			}
			else {
				fwrite(section.data(), sizeof(uint8_t), section.data().length(), file);
			}
		}
	}
	return *this;
}

Linker& Linker::makeFbin(std::string filename)
{
	File file(filename, "wb");
	uint64_t size = 7;
	for (auto& kv : sections.data) {
		for (auto& section : kv.second) {
			size += section.data().length();
		}
	}
	uint16_t cnt = (uint16_t)(size / 512);
	fwrite(&cnt, 2, 1, file);
	fputc(0xbb, file);
	unsigned start = (unsigned)_startAddr;
	fwrite(&start, 4, 1, file);
	fputc(0xff, file);
	fputc(0xe3, file);
	for (auto& kv : sections.data) {
		for (auto& section : kv.second) {
			if (section.sectionType() == SECTION_BSS) {
				for (uint64_t i = 0; i < section.data().length(); i++) {
					fputc(0, file);
				}
			}
			else {
				fwrite(section.data(), sizeof(uint8_t), section.data().length(), file);
			}
		}
	}
	return *this;
}
